import pandas as pd
import numpy as np
from scipy.sparse import csr_matrix

# 读取数据
data = pd.read_csv('ratings.csv', delimiter='\s+')

# 创建用户索引
user_idx = data['userId'].unique()
num_users = len(user_idx)

movie_idx = data['movieId'].unique()
num_movies = len(movie_idx)

# 创建稀疏矩阵
row = data['userId'].map(lambda x: np.where(user_idx == x)[0][0])
col = data['movieId'] - 1
values = data['rating']





# 创建偏置项
global_mean = np.mean(values)

user_bias = np.zeros((num_users,1))
item_bias = np.zeros((num_movies,1))
# 创建随机初始矩阵
user_matrix = np.random.rand(num_users, 5)
item_matrix = np.random.rand(num_movies, 5)
# 迭代训练
learning_rate = 0.000001
num_iterations = 1
lambda_ = 0.01  # 设置 L2 正则化参数
for i in range(num_iterations):

    pred = np.dot(user_matrix, item_matrix.T)+np.mean(values)+user_bias[:,np.newaxis]+item_bias[:,np.newaxis].T
    # 计算预测评分和实际评分之差


    error = values.values - pred[row,col]
    # 更新参数并添加 L2 正则化
    user_bias[row] += learning_rate * (np.sum(error, axis=1) - 2 * lambda_ * user_bias[row])
    item_bias[col] += learning_rate * (np.sum(error, axis=0) - 2 * lambda_ * item_bias[col])

    user_matrix[row, :] += learning_rate * (
                np.sum(error[:, np.newaxis] * item_matrix[col, :], axis=0) - 2 * lambda_ * user_matrix[row, :])
    item_matrix[col, :] += learning_rate * (
                np.sum(error[:, np.newaxis] * user_matrix[row, :], axis=0) - 2 * lambda_ * item_matrix[col, :])

    # 计算均方误差
    mse = np.mean(error ** 2)
    # if i % 100 == 0:
    print("Iteration:", i, "MSE:", mse)

# 将用户矩阵与用户索引关联起来
user_matrix = pd.DataFrame(user_matrix, index=user_idx)

# 预测评分矩阵
pred_ratings = global_mean + user_bias[:, np.newaxis] + item_bias[:, np.newaxis] + np.dot(user_matrix, item_matrix.T)
print(pred_ratings)